#!/usr/bin/env python
import roslib
roslib.load_manifest('ipm_to_laserscan')
import sys
import rospy
import cv
import time
from std_msgs.msg import String
from sensor_msgs.msg import Image
from cv_bridge import CvBridge, CvBridgeError

class hsv_node:

    def __init__(self):
        rospy.init_node('hsv_node')

        """ Create the cv_bridge object """
        self.bridge = CvBridge()

        """ Subscribe to the raw camera image topic """
        self.image_in_topic = rospy.get_param('~image_in_topic')
        self.image_out_topic = rospy.get_param('~image_out_topic')
        self.meters_per_pixel = float(rospy.get_param('~meters_per_pixel'))
        self.camera_offset = float(rospy.get_param('~camera_offset'))

        self.FACTOR = rospy.get_param('~scale_factor', 6)
        self.UPPER = rospy.get_param('~hsv_upper', 90)
        self.LOWER = rospy.get_param('~hsv_lower', 40)

        self.SPATIAL = float(rospy.get_param('~spatial_radius', 15.0))
        self.COLOR = float(rospy.get_param('~color_radius', 30.0))
        self.storage = cv.CreateMemStorage();
        self.point = None

        #cv.NamedWindow("Shift", 0)
        #cv.NamedWindow("Yellow", 0)
        #cv.NamedWindow("Blue", 0)
        #cv.NamedWindow("Merged", 0)
        #cv.NamedWindow("Before", 0)
        #cv.NamedWindow("After", 0)
        #cv.NamedWindow("H", 0)
        #cv.NamedWindow("V", 0)

        self.image_sub = rospy.Subscriber(self.image_in_topic, Image, self.handle_image, queue_size=1)
        self.image_pub = rospy.Publisher(self.image_out_topic,Image)


    def handle_image(self, data):
        try:
            """ Convert the raw image to OpenCV format """
            src = self.bridge.imgmsg_to_cv(data, "bgr8")
            self.process_image(src)
        except CvBridgeError, e:
            rospy.logwarn(e)
            return

    def process_image(self, img):
            self.orig_size = (width, height) = cv.GetSize(img)

            self.small_size = (self.orig_size[0]/self.FACTOR,self.orig_size[1]/self.FACTOR)
            self.small_size_mask = (self.small_size[0] + 2,self.small_size[1] + 2)

            dst = cv.CreateImage(self.small_size, cv.IPL_DEPTH_8U, 3)
            hsv = cv.CreateImage(self.small_size, cv.IPL_DEPTH_8U, 3)
            sm = cv.CreateImage(self.small_size, cv.IPL_DEPTH_8U, 3);
            tmp = cv.CreateImage(self.small_size, cv.IPL_DEPTH_8U, 3);
            
            out = cv.CreateImage(self.orig_size, cv.IPL_DEPTH_8U, 1)
            contour_img = cv.CreateImage(self.orig_size, cv.IPL_DEPTH_8U, 1)

            h_plane = cv.CreateImage(cv.GetSize(sm), cv.IPL_DEPTH_8U, 1);
            v_plane = cv.CreateImage(cv.GetSize(sm), cv.IPL_DEPTH_8U, 1);

            green = cv.CreateImage(cv.GetSize(sm), cv.IPL_DEPTH_8U, 1);
            yellow = cv.CreateImage(cv.GetSize(sm), cv.IPL_DEPTH_8U, 1);
            blue = cv.CreateImage(cv.GetSize(sm), cv.IPL_DEPTH_8U, 1);
            

            cv.Resize(img, sm)

            #cv.PyrMeanShiftFiltering(sm, dst, self.SPATIAL, self.COLOR)
            cv.Smooth(sm, dst, cv.CV_GAUSSIAN)
            #cv.Dilate(dst, dst, iterations =1)
            cv.Smooth(dst, dst, cv.CV_GAUSSIAN)


            cv.CvtColor(dst, hsv, cv.CV_BGR2HSV);
            cv.Split(hsv, h_plane, None, v_plane, None)
            # To be more probablistic, we could calcualte a histogram of a safe region
            # and treshhold on the most largest and smallest H-values that appear most frequently.
            # But this is likely good enough.
            cv.Threshold(h_plane, yellow, 32, 255, cv.CV_THRESH_BINARY_INV)
            cv.Threshold(yellow, yellow, 28, 255, cv.CV_THRESH_TOZERO)
            cv.Dilate(yellow, yellow, iterations =1)

            cv.Threshold(h_plane, blue, 55, 255, cv.CV_THRESH_BINARY_INV)
            cv.Threshold(blue, blue, 50, 255, cv.CV_THRESH_TOZERO)
            # cv.Not(blue, blue)
            #cv.Dilate(blue, blue, iterations =1)

            cv.Threshold(h_plane, h_plane, self.UPPER, 255, cv.CV_THRESH_BINARY_INV)
            cv.Threshold(h_plane, h_plane, self.LOWER, 255, cv.CV_THRESH_TOZERO)
            #cv.Dilate(h_plane, h_plane, iterations =1)

            cv.Threshold(v_plane, v_plane, 25, 255, cv.CV_THRESH_TOZERO)

            #cv.Smooth(v_plane, v_plane)
            #cv.Not(v_plane, v_plane)

            #cv.Or(yellow, v_plane, v_plane)
            #cv.Or(blue, h_plane, h_plane)
            #cv.And(h_plane, v_plane, h_plane)
            #cv.Dilate(h_plane, h_plane, iterations = 1)


            # A gaussian blur may also improve results
            # cv.ShowImage("V", v_plane)

            #publish Image
            cv.Resize(h_plane, out)

            p1 = (0, img.height - (img.height-int(self.camera_offset/self.meters_per_pixel)))
            p2 = (img.width, img.height)
            #cv.Rectangle(out, p1, p2, (0,255,50,20), -1)
            #cv.Dilate(out, out, iterations = 1)
    
            image_message = self.bridge.cv_to_imgmsg(v_plane, encoding="passthrough")
            self.image_pub.publish(image_message)

def main(args):
      vn = hsv_node()
      try:
        rospy.spin()
      except KeyboardInterrupt:
        rospy.logwarn( "Shutting down hsv node.")
      cv.DestroyAllWindows()

if __name__ == '__main__':
    main(sys.argv)

